#define K_MAX_SHAPE_DIM 0
#include "kernel_operator.h"
#include <type_traits>
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
constexpr int32_t SPACE_NUM = 8000;

template<typename TYPE_Y> class KernelGcd_int32 {
    using T = TYPE_Y;
public:
    __aicore__ inline KernelGcd_int32() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
                            uint32_t formerNum, uint32_t tailLength, TPipe* pipeIn) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        uint32_t size = tailLength + 128;
        uint32_t offset = 0;
        if (GetBlockIdx() < formerNum)
        {
            offset = size * GetBlockIdx();
        }
        else
        {
            offset = tailLength * GetBlockIdx() + 128 * formerNum;
            size = tailLength;
        }
        x1Gm.SetGlobalBuffer((__gm__ int32_t*)x1 + offset, size);
        x2Gm.SetGlobalBuffer((__gm__ int32_t*)x2 + offset, size);
        yGm.SetGlobalBuffer((__gm__ int32_t*)y + offset, size);
        tileDataNum = SPACE_NUM;
        loopCount = size/SPACE_NUM;
        tailDataNum = size % SPACE_NUM;
        if(tailDataNum == 0) {tailDataNum = tileDataNum; loopCount-=1;}

        pipeIn->InitBuffer(inQueueX1, BUFFER_NUM, (SPACE_NUM+32)*4);
        pipeIn->InitBuffer(inQueueX2, BUFFER_NUM, (SPACE_NUM+32)*4);
        pipeIn->InitBuffer(outQueueY, BUFFER_NUM, (SPACE_NUM+32)*4);
    }
    
    __aicore__ inline void Process() {
        this->processDataNum = this->tileDataNum;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
        this->processDataNum = this->tailDataNum;
        CopyIn(loopCount);
        Compute(loopCount);
        CopyOut(loopCount);
    }
    
    private:
    __aicore__ inline void CopyIn(int32_t progress) {
        LocalTensor<int32_t> x1Local = inQueueX1.AllocTensor<int32_t>();
        LocalTensor<int32_t> x2Local = inQueueX2.AllocTensor<int32_t>();
        DataCopy(x1Local, x1Gm[progress * this->tileDataNum], this->processDataNum);
        DataCopy(x2Local, x2Gm[progress * this->tileDataNum], this->processDataNum);
        inQueueX1.EnQue(x1Local);
        inQueueX2.EnQue(x2Local);
    }
    __aicore__ inline void Compute(int32_t progress) {
        
        LocalTensor<int32_t> x1Local = inQueueX1.DeQue<int32_t>();
        LocalTensor<int32_t> x2Local = inQueueX2.DeQue<int32_t>();
        LocalTensor<int32_t> yLocal = outQueueY.AllocTensor<int32_t>();
        Cast(x1Local.template ReinterpretCast<float>(), x1Local, RoundMode::CAST_NONE, this->processDataNum);
        Cast(x2Local.template ReinterpretCast<float>(), x2Local, RoundMode::CAST_NONE, this->processDataNum);
        Abs(x1Local.template ReinterpretCast<float>(), x1Local.template ReinterpretCast<float>(),this->processDataNum);
        Abs(x2Local.template ReinterpretCast<float>(), x2Local.template ReinterpretCast<float>(),this->processDataNum);
        
        Max(yLocal.template ReinterpretCast<float>(), x1Local.template ReinterpretCast<float>(), x2Local.template ReinterpretCast<float>(), this->processDataNum);
        WholeReduceMax(yLocal.template ReinterpretCast<float>(), yLocal.template ReinterpretCast<float>(), 64, 125, 1, 1, 8, ReduceOrder::ORDER_ONLY_VALUE);
        ReduceMax(yLocal.template ReinterpretCast<float>(), yLocal.template ReinterpretCast<float>(), yLocal.template ReinterpretCast<float>(), 125, false);
        int32_t loop = ((int32_t)(yLocal.template ReinterpretCast<float>().GetValue(0)))/2;

        Cast(x1Local, x1Local.template ReinterpretCast<float>(), RoundMode::CAST_RINT, this->processDataNum);
        Cast(x2Local, x2Local.template ReinterpretCast<float>(), RoundMode::CAST_RINT, this->processDataNum);

        while(loop --)
        {
            Max(yLocal, x1Local, x2Local, this->processDataNum);
            Min(x1Local, x1Local, x2Local, this->processDataNum);
            Sub(x2Local, yLocal, x1Local, this->processDataNum);
            Max(yLocal, x1Local, x2Local, this->processDataNum);
            Min(x1Local, x1Local, x2Local, this->processDataNum);
            Sub(x2Local, yLocal, x1Local, this->processDataNum);
        }
        Max(yLocal, x1Local, x2Local, this->processDataNum);
        outQueueY.EnQue<int32_t>(yLocal);
        inQueueX1.FreeTensor(x1Local);
        inQueueX2.FreeTensor(x2Local);
    }
    __aicore__ inline void CopyOut(int32_t progress) {
        LocalTensor<int32_t> yLocal = outQueueY.DeQue<int32_t>();
        DataCopy(yGm[progress * this->tileDataNum], yLocal, this->processDataNum);
        outQueueY.FreeTensor(yLocal);
    }
private:
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX1, inQueueX2;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<QuePosition::VECCALC> QueueTmp1, QueueTmpBit, QueueTmpX1, QueueTmpX2, QueueResult, Queuezero;

    GlobalTensor<int32_t> x1Gm;
    GlobalTensor<int32_t> x2Gm;
    GlobalTensor<int32_t> yGm;
    uint32_t loopCount;
    uint32_t tileDataNum;
    uint32_t tailDataNum;
    uint32_t processDataNum;
};
template<typename TYPE_Y> class KernelGcd {
    using T = TYPE_Y;
public:
    __aicore__ inline KernelGcd() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
                            uint32_t size) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->size = size;
        
        x1Gm.SetGlobalBuffer((__gm__ T*)x1, size);
        x2Gm.SetGlobalBuffer((__gm__ T*)x2, size);
        yGm.SetGlobalBuffer((__gm__ T*)y, size);
    }
    
    __aicore__ inline void Process() {
        for(int32_t i=0; i<this->size; i++)
        {
            yGm.SetValue(i, gcd(x1Gm.GetValue(i), x2Gm.GetValue(i)));
        }
    }
    __aicore__ inline T gcd(T m, T n)
    {
        if(n < 0) n = -n;
        if(m < 0) m = -m;
        while(n != 0)
        {
            T temp = m%n;
            m = n;
            n = temp;
        }
        return m;
    }
private:
    TPipe pipe;
    GlobalTensor<DTYPE_X1> x1Gm;
    GlobalTensor<DTYPE_X2> x2Gm;
    GlobalTensor<DTYPE_Y> yGm;
    uint32_t size;
};
template<typename TYPE_Y> class KernelGcd_broadcast {
    using T = TYPE_Y;
public:
    __aicore__ inline KernelGcd_broadcast() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
                            uint32_t size,uint32_t y_dim, uint32_t* x1_shape, uint32_t* x2_shape, uint32_t* y_shape) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->size = size;
        this->y_dim = y_dim;
        for(uint32_t i=0; i<this->y_dim; i++)
        {
            this->x1_shape[i] = x1_shape[i];
            this->x2_shape[i] = x2_shape[i];
            this->y_shape[i] = y_shape[i];
        }
        
        x1Gm.SetGlobalBuffer((__gm__ T*)x1, size);
        x2Gm.SetGlobalBuffer((__gm__ T*)x2, size);
        yGm.SetGlobalBuffer((__gm__ T*)y, size);
    }
    
    __aicore__ inline void Process() {
        for(int32_t i=0; i<this->size; i++)
        {
            yGm.SetValue(i, gcd(x1Gm.GetValue(get_broadcasted_index(i, x1_shape, y_shape, y_dim)), x2Gm.GetValue(get_broadcasted_index(i, x2_shape, y_shape, y_dim))));
        }
    }
    __aicore__ inline T gcd(T m, T n)
    {
        if(n < 0) n = -n;
        if(m < 0) m = -m;
        while(n != 0)
        {
            T temp = m%n;
            m = n;
            n = temp;
        }
        return m;
    }
    __aicore__ inline uint32_t get_broadcasted_index(uint32_t indices, uint32_t *original_shape, uint32_t *broadcast_shape, uint32_t ndim)
    {
        uint32_t shape[8];
        uint32_t original_indices = 0;
        for(int32_t i=ndim-1; i>=0; i--)
        {
            shape[i] = 0;
            if(original_shape[i]!=1)
            {
                shape[i] = indices % broadcast_shape[i];
            }
            indices /= broadcast_shape[i];        
        }
        for(uint32_t i=0; i<ndim; i++)
        {
            original_indices = shape[i] + original_indices * original_shape[i];
        }
        return original_indices;
    }
private:
    TPipe pipe;
    GlobalTensor<DTYPE_X1> x1Gm;
    GlobalTensor<DTYPE_X2> x2Gm;
    GlobalTensor<DTYPE_Y> yGm;
    uint32_t size;

    uint32_t y_dim;
    uint32_t x1_shape[5];
    uint32_t x2_shape[5];
    uint32_t y_shape[5];
};
extern "C" __global__ __aicore__ void gcd(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    if (TILING_KEY_IS(1)) {
        TPipe pipe;
        KernelGcd_int32<int32_t> op;
        op.Init(x1, x2, y, tiling_data.formerNum, tiling_data.tailLength, &pipe);
        op.Process();
    }
    if (TILING_KEY_IS(2)) {
        KernelGcd<DTYPE_X1> op;
        op.Init(x1, x2, y, tiling_data.size);
        op.Process();
    }
    else if (TILING_KEY_IS(3)) {
        KernelGcd_broadcast<DTYPE_X1> op;
        op.Init(x1, x2, y, 
                tiling_data.size, tiling_data.y_dim, tiling_data.x1_shape, tiling_data.x2_shape, tiling_data.y_shape);
        op.Process();
    }
}